package org.sample.mqtt.component.support;

import io.protostuff.LinkedBuffer;
import io.protostuff.ProtostuffIOUtil;
import io.protostuff.Schema;
import io.protostuff.runtime.RuntimeSchema;
import org.sample.mqtt.component.annotations.Topic;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider;
import org.springframework.core.type.filter.AnnotationTypeFilter;
import org.springframework.integration.mapping.BytesMessageMapper;
import org.springframework.integration.mqtt.support.MqttHeaders;
import org.springframework.integration.support.MessageBuilder;
import org.springframework.messaging.Message;

import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * 使用字节传输报文(基于Protostuff)
 * Created by Alan on 2019/8/7
 */
public class MqttProtostuffMessageMapper implements BytesMessageMapper {

    private static final Logger log = LoggerFactory.getLogger(BytesMessageMapper.class);

    private static final ThreadLocal<LinkedBuffer> BUFFER_THREAD_LOCAL = ThreadLocal.withInitial(() -> LinkedBuffer.allocate(512));

    private static final Map<String, Schema> topicMapper = new HashMap<>();
    private static final Map<String, Schema> classMapper = new HashMap<>();

    public MqttProtostuffMessageMapper(String modelPackages) {
        ClassPathScanningCandidateComponentProvider scanner = new ClassPathScanningCandidateComponentProvider(false);
        scanner.addIncludeFilter(new AnnotationTypeFilter(Topic.class));

        String[] packages = modelPackages.split(",");
        for (String aPackage : packages) {
            Set<BeanDefinition> beanDefinitions = scanner.findCandidateComponents(aPackage);
            for (BeanDefinition definition : beanDefinitions) {
                try {
                    addClass(Class.forName(definition.getBeanClassName()));
                } catch (ClassNotFoundException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }

    public static void addClass(Class<?> clazz) {
        Topic annotation = clazz.getAnnotation(Topic.class);
        String value = annotation.value();
        Schema schema = RuntimeSchema.getSchema(clazz);
        topicMapper.put(value.substring(0, value.lastIndexOf("/")), schema);
        classMapper.put(clazz.getName(), schema);
    }

    @Override
    public Message<?> toMessage(byte[] bytes, Map<String, Object> headers) {
        String topic = (String) headers.get(MqttHeaders.TOPIC);
        String action = topic.substring(0, topic.lastIndexOf("/"));

        Schema schema = topicMapper.get(action);
        Object payload;

        if (schema != null) {
            try {
                payload = schema.newMessage();
                ProtostuffIOUtil.mergeFrom(bytes, payload, schema);
            } catch (Exception e) {
                StringBuilder sb = new StringBuilder("\n反序列化失败：");
                sb.append(topic);
                sb.append("\nstring ").append(new String(bytes));
                sb.append("\nbase64 ").append(Base64.getEncoder().encodeToString(bytes));
                log.error(sb.toString(), e);
                throw e;
            }
        } else {
            payload = bytes;
        }
        MessageBuilder<Object> messageBuilder = MessageBuilder.withPayload(payload).copyHeaders(headers);
        return messageBuilder.build();
    }

    @Override
    public byte[] fromMessage(Message<?> message) {
        Object payload = message.getPayload();
        Schema schema = classMapper.get(payload.getClass().getName());
        LinkedBuffer buffer = BUFFER_THREAD_LOCAL.get();
        try {
            byte[] bytes = ProtostuffIOUtil.toByteArray(payload, schema, buffer);
            return bytes;
        } catch (Exception e) {
            throw e;
        } finally {
            buffer.clear();
        }
    }
}